import os
import rospy
import rosbag
import tf2_ros
# from tf2_msgs.msg import TFMessage
from geometry_msgs.msg import Pose, TransformStamped
from process_json_map import process_time_json, output_json

# Replace these with your actual bag file, frames, and timestamp
bag_file = "../bag/test.bag"
source_frame = "base_link"
target_frame = "map"
# Get the odom topic
odom_topic = "/robot_pose_ekf/odom_combined"
image_topic = "/camera/color/image_raw"


def get_first_timestamp(bag_file, topic):
    first_timestamp = None
    with rosbag.Bag(bag_file, "r") as bag:
        for topic, msg, timestamp in bag.read_messages(topics=[topic]):
            if first_timestamp is None:
                first_timestamp = timestamp
            break
    return first_timestamp


def convert_tf_to_dict(transform):
    # 将Pose对象转换为字典
    pose_dict = {
        "position": {"x": transform.translation.x, "y": transform.translation.y, "z": transform.translation.z},
        "orientation": {
            "w": transform.rotation.w,
            "x": transform.rotation.x,
            "y": transform.rotation.y,
            "z": transform.rotation.z,
        },
    }
    return pose_dict


def convert_pose_to_dict(pose):
    # 将Pose对象转换为字典
    pose_dict = {
        "position": {"x": pose.position.x, "y": pose.position.y, "z": pose.position.z},
        "orientation": {
            "w": pose.orientation.w,
            "x": pose.orientation.x,
            "y": pose.orientation.y,
            "z": pose.orientation.z,
        },
    }
    return pose_dict


def convert_dict_to_pose(dict):
    # 将Pose对象转换为字典
    if "position" in dict and "orientation" in dict:
        if (
            "x" in dict["position"]
            and "y" in dict["position"]
            and "z" in dict["position"]
            and "w" in dict["orientation"]
        ):
            pose = Pose()
            pose.position.x = float(dict["position"]["x"])
            pose.position.y = float(dict["position"]["y"])
            pose.position.z = float(dict["position"]["z"])
            pose.orientation.w = float(dict["orientation"]["w"])
            pose.orientation.x = float(dict["orientation"]["x"])
            pose.orientation.y = float(dict["orientation"]["y"])
            pose.orientation.z = float(dict["orientation"]["z"])
        return pose
    return None


def query_transform(tf_buffer, timestamp):
    try:
        transform_stamped = tf_buffer.lookup_transform(target_frame, source_frame, timestamp)
        print(transform_stamped)
        return transform_stamped.transform
    except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
        print(f"Error looking up transform: {e}")
        return None


def query_pose(bag_file, timestamp):
    """Queries the pose from a ROSBag at a given timestamp."""

    with rosbag.Bag(bag_file, "r") as bag:
        # Create a view of the topic with start and end time
        # We set the start and end time to be very close to the desired timestamp
        # to ensure we get the closest message
        start_time = timestamp - rospy.Duration(0.1)
        end_time = timestamp + rospy.Duration(0.1)
        view = bag.read_messages(topics=odom_topic, start_time=start_time, end_time=end_time)
        # Iterate through the messages and find the closest one to the desired timestamp
        closest_message = None
        min_time_diff = float("inf")
        for topic, msg, msg_time in view:
            time_diff = abs(msg_time - timestamp).to_sec()
            if time_diff < min_time_diff:
                min_time_diff = time_diff
                closest_message = msg
                print("timestamp: {} ;  time_diff: {}".format(timestamp.to_sec(), time_diff))

        if closest_message:
            # Extract the pose information from the message
            pose = closest_message.pose.pose
            return pose
        else:
            return None


def main_pose():
    map_dict = process_json(os.path.join(os.path.dirname(__file__), "time.json"))
    first_timestamp = get_first_timestamp(bag_file, image_topic)
    for scene in map_dict:
        for obj in map_dict[scene]:
            print("object name: {}".format(obj))
            pose = query_pose(
                bag_file, first_timestamp + rospy.Duration(map_dict[scene][obj])
            )
            map_dict[scene][obj] = convert_pose_to_dict(pose)
    output_json(map_dict, os.path.join(os.path.dirname(__file__), "map.json"))
    
            
def main_tf():
    map_dict = process_json(os.path.join(os.path.dirname(__file__), "time.json"))
    # Create a tf2 Buffer and TransformListener
    tf_buffer = tf2_ros.Buffer(cache_time=rospy.Duration(500.0))
    first_timestamp = get_first_timestamp(bag_file, image_topic)
    # Open the bag file
    with rosbag.Bag(bag_file, "r") as bag:
        # Iterate through the messages in the bag
        for topic, msg, time in bag.read_messages(topics=['/tf', '/tf_static']):
            for transform in msg.transforms:
                # Add each transform to the buffer
                transform_stamped = TransformStamped()
                transform_stamped.header = transform.header
                transform_stamped.child_frame_id = transform.child_frame_id
                transform_stamped.transform.translation.x = transform.transform.translation.x
                transform_stamped.transform.translation.y = transform.transform.translation.y
                transform_stamped.transform.translation.z = transform.transform.translation.z
                transform_stamped.transform.rotation.x = transform.transform.rotation.x
                transform_stamped.transform.rotation.y = transform.transform.rotation.y
                transform_stamped.transform.rotation.z = transform.transform.rotation.z
                transform_stamped.transform.rotation.w = transform.transform.rotation.w
                tf_buffer.set_transform(transform_stamped, 'default_authority')
    if first_timestamp is None:
        raise ValueError("No transforms found in bag file.")

    # Query the transform
    for scene in map_dict:
        for obj in map_dict[scene]:
            print("\nobject name: {}".format(obj))
            timestamp = first_timestamp + rospy.Duration(map_dict[scene][obj])
            transform = query_transform(tf_buffer, timestamp)
            if transform is not None:
                map_dict[scene][obj] = convert_tf_to_dict(transform)
    output_json(map_dict, os.path.join(os.path.dirname(__file__), "map.json"))
    

if __name__ == "__main__":
    flag_tf = True
    if flag_tf:
        main_tf()
    else:
        main_pose()